require 'lfs'
require 'unity'
local mpi = require "luaMPI"
local nl = require"luaNumeric"
-- local DataSheet = require("DataSheet")


----------------------------------------------------------------------------------------------------------------------------------------
--初始化参数
local PSO={
	swarm_size = 30,           --种群大小
	iterations = 50,		     --最大迭代次数
	inertia = 1.0,                  --惯性权重omega， 较大值代表全局搜索能力更强，较小值适合局部深度搜索
	correction_factor = 2.0,   --学习因子：自我认知和社会认知
	max_val = 1e20,             --最大评价值
	par = { {name='a', pmax=300,pmin=-300}};   --参数空间
	max_vel = 0.2,               --速度限制
	num_iter = 0,	              --当前迭代次数
	cache_data = {},                         
	swarm = {},                               --粒子实体
	results_path = 'iters.txt',             --每一个粒子，每一步迭代过程的数据
	best_results_path = 'best.txt',     --全局最优解存储路径
	cache_path = 'cache.txt',            --缓存目录，
	global_best_value =  1e20,
	db_name = "pso.db",
	db_keys = { "id_swarm", "id_iter", "evaluation" }
};
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--创建表和创建元表
function PSO:new (o)
	o = o or {};
	self.__index = self;
	setmetatable(o,self)
	return o;
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--初始化函数
function PSO:init()
	math.randomseed(0);                        --初始化随机数种子，保证不同时间调用，产生相同结果
	self:init_swarm_array();                    --初始化粒子
	self:read_cache();                             --读取缓存数据
	os.remove( self.best_results_path );   --清除最优粒子目录
        self:writeGlobalBestSwarmFileHead(); --写最优粒子输出文件头
	os.rename( self.results_path, self.results_path..'.back.'..os.date('%Y-%m-%d(%H-%M-%S)'))    --备份计算数据
--~ 	print( lfs.rmdir( self.best_results_path ) );

	-- self.db = DataSheet:new();
	-- self.db:init( self.db_name );
	-- self.db:add_keys(self.db_keys );
	-- self.db:create_data_table();
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--随机函数：产生（pmin，pmax）范围内的随机数
function random_pos( d )
     return d.pmin+(d.pmax-d.pmin)*math.random()
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--- 惯量函数: 随机惯量
function PSO:inertiaFun()
	return math.random() * self.inertia;
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--打印函数
function PSO:print( )
	print( "swarm_size:"..self.swarm_size);
	print( "iterations:"..self.iterations);
	print( "inertia:"..self.inertia);
	print( "correction_factor:"..self.correction_factor);
	print( "max_val:"..self.max_val);
	print( "num_iter:"..self.num_iter);
	print( "swarm arr:"..#self.swarm)
	-- print( "init_type:"..#self.init_type)
	print( "parameter size:"..#self.par)
	print( "parameters:")
	print_table_table(self.par)
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--打印粒子数组
function PSO:print_swarm_array(num)
	if type(num) == 'number' and num > 1 and num < self.swarm_size then		
		print('index:'..num..':')
		    for k,v in pairs(self.par) do
			print( '  k:'..k..' '..v.name..':' )
			print( "    pos:"..self.swarm[num][k].pos);
		end
	else
		for ii=1,self.swarm_size do
		    print('index:'..ii..':')
		    for k,v in pairs(self.par) do
			print( '  k:'..k..' '..v.name..':' )
			print( '    pos:'..self.swarm[ii][k].pos);
		    end
		end
	end
end
------------------------------------------------------------------------------------------------

------------------------------------------------------------------------------------------------
--打印粒子数组
function PSO:print_partical( partical )
    local str='';
    for k,v in pairs(self.par) do
        str = str..v.name..'='..tostring(partical[k].best_pos)..'  ';
    end
    str = str..partical.best_value
    print(str);
end
------------------------------------------------------------------------------------------------


-----------------------------------------------------------------------------
--读取文件缓存
function PSO:read_cache( path )
	local file = io.open( self.cache_path,"r");
	if file==nil then return; end
	assert( file, "file:"..self.cache_path.." open fail!");
	for line in file:lines() do
--~ 	     print( line )
	     vals = string.split( line, "\t" ) 
	    vals = table2number( vals )
--~ 	     print_table_val( vals );
	     table.insert(self.cache_data, vals)
        end
	file:close()    
	print( "read line:", #self.cache_data );
	-- print_table_table( self.cache_data );
end
-----------------------------------------------------------------------------


-----------------------------------------------------------------------------
--打包粒子数据为表
function PSO:dump_particle( particle )
     local data={};
      data[1] = self.num_iter;
      data[2] = particle.index
      for k,v in pairs(self.par) do
           data[2+k] = particle[k].pos;	
	   data[2+k + #self.par] = particle[k].vel; 
	   data[2+k + 2*(#self.par)] = particle[k].best_pos; 
     end
     return data;
end
-----------------------------------------------------------------------------


-----------------------------------------------------------------------------
--查找缓存是否存在已经计算完成的数据，如果存在，返回data
--如果不存在，返回nil
function PSO:find_cache( particle )
     data = self:dump_particle( particle )
     --print_table(data)
     
     for k,v in pairs( self.cache_data ) do
           local flag = true;
	   for kd,vk in pairs( data ) do 	       
		if math.abs( vk - (v[kd]) ) > 1e-8 then
		     flag = false;		    
		end
	   end
	   if flag then
--~ 	        print_table( v )
		return v;
	   end
     end
     return nil;
end
-----------------------------------------------------------------------------

-----------------------------------------------------------------------------
--write the calculation data of partical to file
function PSO:write2file( particle )
--~ 	print( self.results_path );
	local file = io.open(self.results_path, "a")
	file:write(self.num_iter,"\t", particle.index)
	for k,v in pairs(self.par) do
             file:write("\t",particle[k].pos)
        end
	for k,v in pairs(self.par) do
             file:write("\t",particle[k].vel)
        end
	for k,v in pairs(self.par) do
             file:write("\t",particle[k].best_pos)
        end
        file:write("\t",particle.value,"\n")
	file:close()	
end
-----------------------------------------------------------------------------

------------------------------------------------------------------------------------------------
--write the head of best partical  file
function PSO:writeGlobalBestSwarmFileHead()
	local file = io.open(self.best_results_path, "w")
	file:write( "iterate", "\t", "best_swarm_index","\t", "best_num_iter");
	for k,v in pairs(self.par) do
             file:write("\t",v.name)
        end
	file:write("\t","global_best_value","\n")
	file:close()	
end
------------------------------------------------------------------------------------------------

------------------------------------------------------------------------------------------------
--write the  best partical  info to file
function PSO:writebest2file( )
	local file = io.open(self.best_results_path, "a")
	--file:write(self.num_iter,"\t", self.best_swarm.index,"\t")
	file:write(self.num_iter,"\t", self.best_swarm.index,"\t",self.best_swarm.best_num_iter)
	for k,v in pairs(self.par) do
             file:write("\t",self.best_swarm[k].best_pos)
        end
	file:write("\t",self.global_best_value,"\n")
	file:close()	
end
------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--初始随机
function PSO:initVelocityFun(id_par)	
	return 0;   --初始化速度位置
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
--初始化粒子数组
function PSO:init_swarm_array()	
	self.swarm={};
	for ii=1,self.swarm_size do
	    self.swarm[ii]={}
	    for k,v in pairs(self.par) do
		self.swarm[ii][k] = {}
		self.swarm[ii][k].pos = random_pos(v);  --初始化位置		
		self.swarm[ii][k].vel = self:initVelocityFun(k);  --初始化速度 initial velocity
		self.swarm[ii][k].best_pos = self.swarm[ii][k].pos;  --初始化最优位置	
	    end
	    self.swarm[ii].best_value = self.max_val;  --初始化适用度            self.swarm[ii].best_num_iter = 1;  --初始化最优粒子编号
	    self.swarm[ii].index = ii;  --粒子编号
	end	
	self:init_samples_swarm_array(); --初始化样本数组
end
----------------------------------------------------------------------------------------------------------------------------------------




----------------------------------------------------------------------------------------------------------------------------------------
--初始随机
function PSO:initSampleVelocityFun(id_par,vel)	
	return vel or 0;  --初始化速度 initial velocity
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--初始化样本数组
function PSO:init_samples_swarm_array()	
	if self.samples ~= nil then
		for ii=1,#self.samples do
			for k,v in pairs(self.par) do
				self.swarm[ii][k].pos = self.samples[ii][k].pos;         --初始化位置为样本位置
 				self.swarm[ii][k].vel = self:initSampleVelocityFun(k,self.samples[ii][k].vel);             --初始化位置为速度位置
				self.swarm[ii][k].best_pos = self.swarm[ii][k].pos;    --初始化最优位置	
 			end
		end
	end	
end
----------------------------------------------------------------------------------------------------------------------------------------



----------------------------------------------------------------------------------------------------------------------------------------
--在找出全局最优粒子（比较每一个粒子的历史最优数据）
function PSO:get_global_best_value()
    local min=self.swarm[1].best_value;
    local idx = 1;
    for ii=2,self.swarm_size do
        if min >= self.swarm[ii].best_value then
	    min = self.swarm[ii].best_value;
	    idx = ii;
	end
    end
    return self.swarm[idx],min;
end
----------------------------------------------------------------------------------------------------------------------------------------

function PSO:envaluation( partical )
     local val = self.max_val;
     for k,v in pairs(self.par) do
         if partical[k].pos < v.pmin or partical[k].pos > v.pmax then
	     return self.max_val;
	 end	 
     end
     return math.pow((partical[1].pos - 15),2)
end

---------------------------------------------------------------------------------
--判断粒子位置是否超出给定边界
function PSO:check_boundary( partical )     
     for k,v in pairs(self.par) do
         if partical[k].pos < v.pmin or partical[k].pos > v.pmax then
	     return false;
	 end	 
     end
     return true;
end
---------------------------------------------------------------------------------

---------------------------------------------------------------------------------
--保存粒子到数据库
function PSO:write_db( id_swarm, value )
     self.db:add_data{ id_swarm=id_swarm, id_iter=self.num_iter, evaluation=value }
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--update  position 
function PSO:update_position()	
	mpi.world.barrier();
	for jj = 1,self.swarm_size do
		-- data = std.vector:new( #self.par );
		data = nl.new_vector(#self.par);
		 if mpi.world.rank() == 0 then 			
			for k,v in pairs(self.par) do
				self.swarm[jj][k].pos = self.swarm[jj][k].pos + self.swarm[jj][k].vel /1.3; -- update  position
				--self.swarm[jj][k].pos = self.swarm[jj][k].pos + self.swarm[jj][k].vel; -- update  position
				data[k] = self.swarm[jj][k].pos;
			end		
			mpi.world.broadcast(data,0)
		else
			mpi.world.broadcast(data,0);
			for k,v in pairs(self.par) do
				self.swarm[jj][k].pos = data[k];
			end		
		end	
	end
	mpi.world.barrier();	
end
----------------------------------------------------------------------------------------------------------------------------------------

----------------------------------------------------------------------------------------------------------------------------------------
-- evaluating position & quality ---
function PSO:evaluate()	
	collectgarbage()
	print("Collect Garbage");
	collectgarbage()
	
	mpi.world.barrier();
	local reqs = mpi.request_list:new();
	local vals = {};
	for jj = 1 , self.swarm_size do
		if ( jj % mpi.world.size() ) == mpi.world.rank()  then
			 -- vals[jj] = std.vector:new( 1 );
			 vals[jj] = nl.new_vector({1});
			 local val;
			 --print( "num_iter:",self.num_iter)
			 local cache_res = self:find_cache( self.swarm[jj] );	
--~ 			print("cache_res:"..cache_res)
		         if cache_res ~= nil then
			     val = cache_res[ #cache_res ];			     
		        else
			     --print("cache_res not finded")
  			     val,res  = self:envaluation(self.swarm[jj]);
			     self:write_db(jj, val, res );
			end
			vals[jj][1] = val;	
			if mpi.world.rank() ~= 0 then			
				req = mpi.world.isend(0,jj,vals[jj]);
				reqs:push_back( req );		
--~ 				print( mpi.world.rank(), "send to", 0, "tag:" , jj ,val)				
			end
			
		end
		if(  mpi.world.rank()==0 ) then
			if( jj % mpi.world.size() ~= 0 ) then
				-- vals[jj] = std.vector:new();
				vals[jj] = nl.new_vector();
				req = mpi.world.irecv( jj % mpi.world.size(),  jj, vals[jj] );
				reqs:push_back( req );	
--~ 				print( mpi.world.rank(), "recv from", jj % mpi.world.size(), "tag:", jj  )	
			end
		end
	end
	----------------------------------------------------------------------------------------------------------------------------------------
	
	----------------------------------------------------------------------------------------------------------------------------------------
	--wait for all partical envaluation finish
	mpi.wait_all(reqs);
	mpi.world.barrier();
	----------------------------------------------------------------------------------------------------------------------------------------
	
	
	
	
	----------------------------------------------------------------------------------------------------------------------------------------
	--broadcast values to all mpi process
	for jj = 1 , self.swarm_size do	  
	   if mpi.world.rank() == 0 then	
		mpi.world.broadcast(vals[jj],0);
	   else
	        -- local val = std.vector:new();
	        local val = nl.new_vector();
		mpi.world.broadcast(val,0);
		vals[jj] = val;
	   end
	end
	----------------------------------------------------------------------------------------------------------------------------------------
	
	----------------------------------------------------------------------------------------------------------------------------------------
	-- update val from recv data from mpi process 
        -- and save position and velocity array to file
	for jj = 1 , self.swarm_size do
		self.swarm[jj].value = vals[jj][1]
		if mpi.world.rank() == 0 then	
			self:write2file( self.swarm[jj] );
		end
	end
	----------------------------------------------------------------------------------------------------------------------------------------

	mpi.world.barrier();
end
----------------------------------------------------------------------------------------------------------------------------------------


function PSO:update_best_pos_and_val()	
	----------------------------------------------------------------------------------------------------------------------------------------
	-- update individual and global best position and val
	for jj = 1 , self.swarm_size do			
		if self.swarm[jj].value < self.swarm[jj].best_value then              -- if new position is better
		    for k,v in pairs(self.par) do
			 self.swarm[jj][k].best_pos =  self.swarm[jj][k].pos ; -- update best position
		    end
		    self.swarm[jj].best_value = self.swarm[jj].value;              -- % and best value
		    self.swarm[jj].best_num_iter= self.num_iter;              -- % and the iterations
		end       
	end	    
	best_swarm, self.global_best_value= self:get_global_best_value();	
	self.best_swarm = best_swarm;	
	----------------------------------------------------------------------------------------------------------------------------------------
	
	----------------------------------------------------------------------------------------------------------------------------------------
	-- save best position and val
	if mpi.world.rank() == 0 then	
		self:writebest2file( );
	end
	----------------------------------------------------------------------------------------------------------------------------------------
end

----------------------------------------------------------------------------------------------------------------------------------------
--- updating velocity vectors
function PSO:update_velocity()
	for jj = 1 , self.swarm_size do		
		for k,v in pairs(self.par) do
		     -- self.swarm[jj][k].vel = math.random() * self.inertia * self.swarm[jj][k].vel 
		     --self.swarm[jj][k].vel =      self.inertia * self.swarm[jj][k].vel 
		     --self.swarm[jj][k].vel =    ( self.iterations - self.num_iter ) / self.iterations *  self.inertia * self.swarm[jj][k].vel
		     self.swarm[jj][k].vel =   self:inertiaFun() * self.swarm[jj][k].vel 
							    + self.correction_factor*math.random()*(self.swarm[jj][k].best_pos - self.swarm[jj][k].pos) 
							    + self.correction_factor*math.random()*(best_swarm[k].best_pos - self.swarm[jj][k].pos)
			if self.swarm[jj][k].vel>self.max_vel*(self.par[k].pmax-self.par[k].pmin) then
				self.swarm[jj][k].vel = self.max_vel*(self.par[k].pmax-self.par[k].pmin)
			elseif self.swarm[jj][k].vel<-1*self.max_vel*(self.par[k].pmax-self.par[k].pmin) then
				self.swarm[jj][k].vel = -1*self.max_vel*(self.par[k].pmax-self.par[k].pmin)
			end
		end
	end 	
end
----------------------------------------------------------------------------------------------------------------------------------------


----------------------------------------------------------------------------------------------------------------------------------------
--处理超出边界的粒子
function PSO:boundary_process()
	if self.boundary_type == 1 then
		for jj = 1,self.swarm_size do
			for k,v in pairs(self.par) do
				if self.swarm[jj][k].pos > v.pmax then
					self.swarm[jj][k].pos = v.pmax
				elseif self.swarm[jj][k].pos < v.pmin then
					self.swarm[jj][k].pos = v.pmin
				end
			end
		end
	end
end
----------------------------------------------------------------------------------------------------------------------------------------

function PSO:step()	
	self.num_iter = self.num_iter+1;	 	
	print( "num_iter:",self.num_iter)
		
	----------------------------------------------------------------------------------------------------------------------------------------
	--envaluation 
	 self:evaluate()
	----------------------------------------------------------------------------------------------------------------------------------------

	
	----------------------------------------------------------------------------------------------------------------------------------------
	--update  individual and global best position and val
	self:update_best_pos_and_val()
	----------------------------------------------------------------------------------------------------------------------------------------
	
	----------------------------------------------------------------------------------------------------------------------------------------
	--- updating velocity vectors
	self:update_velocity()
	----------------------------------------------------------------------------------------------------------------------------------------

	----------------------------------------------------------------------------------------------------------------------------------------
	 --update  position 
	 self:update_position()
	 self:boundary_process()
	----------------------------------------------------------------------------------------------------------------------------------------

end


function PSO:iterate()
	for ii=1,self.iterations do	
		print("--------------------------------------------")	
		print("--Iteration start: step"..ii)			
		self:step();
		print("--Finish------------------------------------------\n")	

	end
end


function PSO:run()
	self:init() -- read_cache, init_swarm_array, rm_best_file
	self:iterate()
end

return PSO;



